import logging
from typing import Tuple

import nmb.NetBIOS
from impacket.dcerpc.v5 import epm, nrpc, transport

from common.common_consts.timeouts import LONG_REQUEST_TIMEOUT, MEDIUM_REQUEST_TIMEOUT
from infection_monkey.i_puppet import TargetHost

from .exceptions import DomainControllerNameFetchError

logger = logging.getLogger(__name__)


def connect_to_dc(dc_ip: str) -> object:
    binding = epm.hept_map(dc_ip, nrpc.MSRPC_UUID_NRPC, protocol="ncacn_ip_tcp")
    rpc_transport = transport.DCERPCTransportFactory(binding)
    rpc_transport.set_connect_timeout(LONG_REQUEST_TIMEOUT)
    rpc_con = rpc_transport.get_dce_rpc()
    rpc_con.connect()
    rpc_con.bind(nrpc.MSRPC_UUID_NRPC)
    return rpc_con


def get_dc_details(host: TargetHost) -> Tuple[str, str, str]:
    dc_ip = str(host.ip)
    dc_name = _get_dc_name(dc_ip=dc_ip)
    dc_handle = "\\\\" + dc_name
    return dc_ip, dc_name, dc_handle


def _get_dc_name(dc_ip: str) -> str:
    """
    Gets NetBIOS name of the Domain Controller (DC).
    """
    nb = nmb.NetBIOS.NetBIOS()
    name = nb.queryIPForName(
        ip=dc_ip, timeout=MEDIUM_REQUEST_TIMEOUT
    )  # returns either a list of NetBIOS names or None

    if name:
        return name[0]
    else:
        raise DomainControllerNameFetchError(
            "Couldn't get domain controller's name, maybe it's on external network?"
        )
